# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import json

import luigi
import pandas as pd
from absl import app, flags
from luigi.local_target import LocalTarget
from sqlitedict import SqliteDict

from ops.bitmex_param_autopick.luigi_base import (CommandTask, MarketSimTask, SqliteTarget)
from ops.bitmex_param_autopick.ops import (generate_runner_config_from_portfolio_db)
from coin.strategy.bitmex_param_picker import picker as param_picker


def get_default_run_db():
  return '/remote/iosg/home/jaewon/ops/bitmex/ops_data.db'


class DailyIdempotentTask(luigi.Config):
  date = luigi.DateParameter(default=datetime.datetime.now().date())
  run_db = luigi.Parameter(default=get_default_run_db())


class PullFeedTask(DailyIdempotentTask):
  def requires(self):
    date_str = self.date.strftime('%Y%m%d')
    return CommandTask(run_id='pull-feed-%s' % date_str,
                       run_db=self.run_db,
                       args=['ops/bitmex_param_autopick/steps/pull_feed.sh'])

  def output(self):
    return self.input()


class CopyBitmexUnavailabilityTask(DailyIdempotentTask):
  def requires(self):
    date_str = self.date.strftime('%Y%m%d')
    return CommandTask(run_id='bitmex-unavail-%s' % date_str,
                       run_db=self.run_db,
                       args=['ops/bitmex_param_autopick/steps/copy_bitmex_unavailability.sh'])

  def output(self):
    return self.input()


class DailyMarketSimTask(DailyIdempotentTask):
  def requires(self):
    return [
        PullFeedTask(date=self.date, run_db=self.run_db),
        CopyBitmexUnavailabilityTask(date=self.date, run_db=self.run_db)
    ]

  def _append_sim_job_id(self, job_id):
    prev_date = self.date - datetime.timedelta(days=1)
    prev_date_str = prev_date.strftime('%Y%m%d')
    with SqliteDict(self.run_db,
                    tablename='sim_job_id',
                    autocommit=True,
                    encode=json.dumps,
                    decode=json.loads) as table:
      job_ids = [str(job_id)]
      if prev_date_str in table:
        job_ids = table[prev_date_str] + job_ids
      table[prev_date_str] = job_ids

  def run(self):
    date_str = self.date.strftime('%Y%m%d')
    prev_date = self.date - datetime.timedelta(days=1)
    prev_date_str = prev_date.strftime('%Y%m%d')
    sim_args = {
        'profile_name': 'pass_unhedge_xbtusd_daily',
        'profile_args': ['--from_date', prev_date_str],
        'max_strategy_per_process': 4,
        'tag': 'xbtusd-daily',
        'sbatch_args': '-n 200',
        'flow_root_dir': '/remote/iosg/home/jaewon/data/flow'
    }

    market_sim_task = MarketSimTask(run_id='market-sim-%s' % date_str,
                                    run_db=self.run_db,
                                    sim_args=sim_args)
    yield market_sim_task

    job_id = market_sim_task.output().get()
    self._append_sim_job_id(job_id)
    self.output().touch()

  def output(self):
    date_str = self.date.strftime('%Y%m%d')
    return SqliteTarget(update_id='daily-market-sim-%s' % date_str, db=self.run_db)


class PortfolioTask(DailyIdempotentTask):
  verbose = luigi.BoolParameter(default=True)

  def requires(self):
    return DailyMarketSimTask(date=self.date, run_db=self.run_db)

  def run(self):
    from_date = self.date - datetime.timedelta(days=80)
    to_date = self.date - datetime.timedelta(days=1)
    sim_result_df = param_picker.load_sim_result(from_date=from_date,
                                                 to_date=to_date,
                                                 from_sim_job_id_db=self.run_db)
    portfolio = param_picker.pick_daily_param(sim_result_df,
                                              self.date,
                                              one_side_stack=False,
                                              to_portfolio_db=self.run_db)
    logging.info('Choosen: %s', ', '.join(portfolio))

    if self.verbose:
      for param_name in portfolio:
        print(param_name)

    self.output().touch()

  def output(self):
    date_str = self.date.strftime('%Y%m%d')
    return SqliteTarget(update_id='portfolio-%s' % date_str, db=self.run_db)


class YesterdayPortfolioSummaryTask(DailyIdempotentTask):
  def requires(self):
    return DailyMarketSimTask(date=self.date, run_db=self.run_db)

  def run(self):
    from_date = self.date - datetime.timedelta(days=80)
    to_date = self.date - datetime.timedelta(days=1)
    dump_from_date = self.date - datetime.timedelta(days=20)

    sim_result_df = param_picker.load_sim_result(from_date=from_date,
                                                 to_date=to_date,
                                                 from_sim_job_id_db=self.run_db)

    with pd.option_context('display.max_rows',
                           None,
                           'display.max_columns',
                           None,
                           'display.width',
                           250,
                           'display.float_format',
                           '{:,.2f}'.format):
      summary_str = param_picker.run_reference_mode(sim_result_df, dump_from_date=dump_from_date)
    print('\n\n\n%s' % summary_str)

    self.output().touch()

  def output(self):
    date_str = self.date.strftime('%Y%m%d')
    return SqliteTarget(update_id='yesterday-portfolio-summary-%s' % date_str, db=self.run_db)


class GenerateRunnerConfigTask(DailyIdempotentTask):
  max_pos = luigi.IntParameter(default=3000)
  num_portfolio = luigi.IntParameter(default=11)

  def requires(self):
    return PortfolioTask(date=self.date, run_db=self.run_db)

  def run(self):
    generate_runner_config_from_portfolio_db(self.run_db,
                                             self.date,
                                             self.max_pos,
                                             check_num_portfolio=self.num_portfolio,
                                             dump_yaml_filepath=self.output().path)

  def output(self):
    date_str = self.date.strftime('%Y%m%d')
    config_filepath = ('/remote/iosg/home/jaewon/ops/bitmex/out/bitmex-pass-unhedge-%s.yaml'
                       % date_str)
    return LocalTarget(config_filepath)


class DeployRunnerConfigTask(DailyIdempotentTask):
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)

    self._gen_config_task = GenerateRunnerConfigTask(date=self.date, run_db=self.run_db)

    date_str = self.date.strftime('%Y%m%d')
    self._deploy_config_task = CommandTask(
        run_id='deploy-config-%s' % date_str,
        run_db=self.run_db,
        args=[
            'ops/bitmex_param_autopick/steps/deploy_runner_config.sh',
            self._gen_config_task.output().path
        ])

  def requires(self):
    return self._gen_config_task

  def run(self):
    yield self._deploy_config_task

  def output(self):
    return self._deploy_config_task.output()


class DailyTask(DailyIdempotentTask):
  def requires(self):
    return [
        DeployRunnerConfigTask(date=self.date, run_db=self.run_db),
        YesterdayPortfolioSummaryTask(date=self.date, run_db=self.run_db)
    ]

  def complete(self):
    return all(r.complete() for r in self.requires())


def main(argv):
  FLAGS = flags.FLAGS

  date = datetime.datetime.now().date()
  if FLAGS.date:
    date = datetime.datetime.strptime(FLAGS.date, '%Y%m%d').date()

  tasks = [DailyTask(date=date)]
  luigi.build(tasks, local_scheduler=True)


if __name__ == '__main__':
  logging.getLogger('sqlitedict').setLevel(logging.WARNING)

  flags.DEFINE_string('date', None, 'If specified, run pipeline for the speicifed date.')

  app.run(main)
